import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr

np.set_printoptions(precision=4, suppress=True)


# absorbing state:吸收态是指一旦进入就永远无法离开的状态。数学上，状态 i 是吸收态当且仅当：
# P[i,i] = 1 且 P[i,j] = 0 (对于所有 j ≠ i)

# transient state:瞬时状态是指可能被吸收，也可能在非吸收态之间转移的状态。
# 从瞬时状态出发，存在一定的概率永远不会返回。

# absorbing probability:从某个瞬时状态出发，最终被某个吸收态吸收的概率。


print("Section3: Absorbing states example")
'''
状态0：吸收态（一旦进入就永远停留）
状态1：瞬时状态（可能被吸收或自我循环）
状态2：瞬时状态（可能转移到状态1或自我循环）
'''
P3 = np.array([
    [1.0, 0.0, 0.0], # 状态0: 吸收态（100%停留在自己）
    [0.5, 0.5, 0.0], # 状态1: 50%被吸收，50%停留在自己
    [0.0, 0.4, 0.6] # 状态2: 40%转到状态1，60%停留在自己
])
print(pd.DataFrame(P3))
print("States:0 is absorbing.Simulate many runs from state 1 and estimate absorbption probability")


def absorption_states(P, start, absorb_states, trials=5000, max_steps=1000):
    times = [] # 存储每次被吸收所需的时间
    absorbed_count = 0 # 统计被吸收的次数
    for _ in range(trials):
        cur = start # 从指定起始状态开始
        for t in range(max_steps): #如果达到最大步数仍未吸收，这次试验不计入吸收统计
            cur = npr.choice(a=len(P), p=P[cur])
            if cur in absorb_states: # 如果进入吸收态
                times.append(t + 1) # 记录吸收时间（步数）
                absorbed_count += 1
                break
    return absorbed_count / trials, np.mean(times) if times else None


# 从状态1开始
# 吸收态集合为{0}

prob, avg_time = absorption_states(P=P3, start=1, absorb_states={0}, trials=5000)
'''
f(i)表示从状态i出发最终被状态0吸收的概率
f(0) = 1 (吸收态自身)
f(1) = 0.5 * f(0) + 0.5 * f(1) + 0.0 * f(2)
f(2) = 0.0 * f(0) + 0.4 * f(1) + 0.6 * f(2)

解：
f(0) = 1
f(1) = 1
f(2) = 1
'''
print(f"Estimated absorption prob (to 0) starting from 1: {prob:.4f}, avg_time:{avg_time:.4f}")
